import numpy as np
import matplotlib.pyplot as plt
from scipy import integrate

# 中文和负号的正常显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 圆盘法计算旋转体体积
def disk_method(f, a, b, axis='x'):
    """
    使用圆盘法计算旋转体体积
    f: 函数
    a, b: 积分区间
    axis: 旋转轴，'x'或'y'
    """
    if axis == 'x':
        # 绕x轴旋转，对x积分
        volume, error = integrate.quad(lambda x: np.pi * (f(x))**2, a, b)
    else:
        # 绕y轴旋转，需要反函数关系
        # 这里假设f是x关于y的函数
        volume, error = integrate.quad(lambda y: np.pi * (f(y))**2, a, b)
    
    return volume, error

# 示例1：y = sqrt(x) 在 [0,4] 上绕x轴旋转
def f1(x):
    return np.sqrt(x)

volume1, error1 = disk_method(f1, 0, 4, 'x')
print(f"旋转体体积1: {volume1:.6f} (理论值: {8*np.pi:.6f})")

# 示例2：椭圆绕x轴旋转
a, b = 2, 1  # 椭圆长半轴和短半轴
def ellipse_y(x):
    return (b/a) * np.sqrt(a**2 - x**2)

volume2, error2 = disk_method(ellipse_y, -a, a, 'x')
theoretical_volume = (4/3) * np.pi * a * b**2
print(f"旋转椭球体体积: {volume2:.6f} (理论值: {theoretical_volume:.6f})")

# 可视化
fig = plt.figure(figsize=(12, 5))

# 第一个子图：y = sqrt(x)
ax1 = fig.add_subplot(121, projection='3d')
x1 = np.linspace(0, 4, 100)
theta = np.linspace(0, 2*np.pi, 100)
X1, Theta = np.meshgrid(x1, theta)
Y1 = f1(X1) * np.cos(Theta)
Z1 = f1(X1) * np.sin(Theta)
ax1.plot_surface(X1, Y1, Z1, alpha=0.7, cmap='viridis')
ax1.set_title(r'$y = \sqrt{x}$ 绕x轴旋转')

# 第二个子图：椭圆旋转
ax2 = fig.add_subplot(122, projection='3d')
x2 = np.linspace(-a, a, 100)
X2, Theta2 = np.meshgrid(x2, theta)
Y2 = ellipse_y(X2) * np.cos(Theta2)
Z2 = ellipse_y(X2) * np.sin(Theta2)
ax2.plot_surface(X2, Y2, Z2, alpha=0.7, cmap='plasma')
ax2.set_title('椭圆绕x轴旋转')

plt.tight_layout()
plt.show()